import pandas as pd
import os
import plotly.graph_objects as go

# --- 数据加载和处理函数 (与您的版本相同，保持不变) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """
    根据时间范围加载一个或多个月份的CSV数据。
    """
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)

    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    # 确定需要读取的月份文件
    date_range = pd.date_range(start=start_date, end=end_date, freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0:
        date_range = [start_date.strftime('%Y-%m')]

    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")

    for year_month in date_range:
        file_name = f"{symbol}-{timeframe}-{year_month}.csv"
        file_path = os.path.join(data_folder_path, file_name)
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                df = pd.read_csv(file_path)
                all_dfs.append(df)
            except Exception as e:
                print(f"读取文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    """
    处理原始数据，包括转换时间戳、设置索引和筛选时间范围。
    """
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume'] # VWAP只需要这几列
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)

    # 筛选出最终需要的时间段
    df_filtered = df.loc[start_time_str:end_time_str]
    print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    return df_filtered


# --- 新增：指标计算函数 ---

def calculate_bollinger_bands(df: pd.DataFrame, length: int = 20, std: float = 2.0) -> pd.DataFrame:
    """
    计算布林带 (BBands)。
    :param df: 输入的DataFrame，需要有 'close' 列。
    :param length: 移动平均线的周期。
    :param std: 标准差的倍数。
    :return: 带有布林带上、中、下轨三列的DataFrame。
    """
    df['bband_middle'] = df['close'].rolling(window=length).mean()
    df['bband_std'] = df['close'].rolling(window=length).std()
    df['bband_upper'] = df['bband_middle'] + (df['bband_std'] * std)
    df['bband_lower'] = df['bband_middle'] - (df['bband_std'] * std)
    print(f"布林带计算完成 (周期={length}, 标准差={std})。")
    return df


def calculate_vwap(df: pd.DataFrame) -> pd.DataFrame:
    """
    计算成交量加权平均价 (VWAP)。
    由于数据是按天加载的，这里的 cumsum() 自然实现了每日重置的效果。
    """
    typical_price = (df['high'] + df['low'] + df['close']) / 3
    tpv = typical_price * df['volume']

    df['vwap'] = tpv.cumsum() / df['volume'].cumsum()
    print("VWAP 计算完成。")
    return df


# --- 新增：绘图函数 ---

def plot_price_with_indicators(df: pd.DataFrame, chart_title: str, bband_length: int, bband_std: float):
    """
    绘制收盘价、布林带和VWAP的折线图。
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return

    fig = go.Figure()

    # 为了实现填充效果，先画下轨，再画上轨
    # 1. 布林带下轨
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['bband_lower'],
        mode='lines',
        line=dict(width=0.5, color='gray'),
        legendgroup='bollinger',
        name=f'BBands Lower ({bband_length}, {bband_std})'
    ))

    # 2. 布林带上轨 (并填充与下轨之间的区域)
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['bband_upper'],
        mode='lines',
        line=dict(width=0.5, color='gray'),
        fill='tonexty',  # 填充到上一个 trace (即下轨)
        fillcolor='rgba(128, 128, 128, 0.2)',
        legendgroup='bollinger',
        name=f'BBands Upper ({bband_length}, {bband_std})'
    ))

    # 3. 布林带中轨
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['bband_middle'],
        mode='lines',
        line=dict(width=1, color='gray', dash='dash'),
        name=f'BBands Middle ({bband_length}, {bband_std})'
    ))

    # 4. 收盘价
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['close'],
        mode='lines',
        line=dict(color='blue', width=2),
        name='Close Price'
    ))

    # 5. VWAP
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['vwap'],
        mode='lines',
        line=dict(color='orange', width=2, dash='dash'),
        name='VWAP'
    ))

    # 更新图表布局
    fig.update_layout(
        title=chart_title,
        xaxis_title='Time (UTC)',
        yaxis_title='Price (USDT)',
        xaxis_rangeslider_visible=False,
        legend_title="Indicators",
        hovermode='x unified' # 统一显示X轴上所有数据点的信息
    )

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "ETHUSDT"
    TIMEFRAME = "5m"

    # --- 核心配置 ---
    # 您想分析的日期 (格式: YYYY-MM-DD)
    TARGET_DATE = "2025-06-13"

    # 布林带参数
    BBAND_LENGTH = 20  # 周期
    BBAND_STD = 2.0    # 标准差倍数
    # ====================================================

    print("--- K线指标图表程序启动 ---")

    # 根据目标日期自动生成开始和结束时间
    start_time_str = f"{TARGET_DATE} 00:00:00"
    end_time_str = f"{TARGET_DATE} 23:59:59"

    # 1. 加载原始数据
    raw_data_df = load_data_in_range(start_time_str, end_time_str, BASE_DATA_PATH, SYMBOL, TIMEFRAME)

    # 2. 预处理和筛选数据
    df_day = process_and_filter_data(raw_data_df, start_time_str, end_time_str)

    if not df_day.empty:
        # 3. 计算技术指标
        df_day = calculate_bollinger_bands(df_day, length=BBAND_LENGTH, std=BBAND_STD)
        df_day = calculate_vwap(df_day)

        # 4. 绘图
        chart_title_str = f"{SYMBOL} {TIMEFRAME} Chart on {TARGET_DATE} (UTC)"
        plot_price_with_indicators(df_day, chart_title=chart_title_str, bband_length=BBAND_LENGTH, bband_std=BBAND_STD)

        print("--- 图表生成完毕 ---")
    else:
        print(f"\n在日期 {TARGET_DATE} 未能加载或处理任何数据，程序退出。请检查您的配置和文件路径。")